"""
@author: j.l.liu
@create: 
@description: LangChain 把需求文本拆字段 → DeepSeek 按 Gherkin 语法生成测试场景
"""
import os
from pathlib import Path

from fastapi import FastAPI, HTTPException, UploadFile, File
from langchain.chat_models import ChatOpenAI
from langchain.document_loaders import TextLoader
from langchain.prompts import ChatPromptTemplate
from langchain.text_splitter import RecursiveCharacterTextSplitter
from starlette.responses import PlainTextResponse

app = FastAPI(title="Auto-Gherkin Generator")
UPLOAD_DIR = Path("uploads")
UPLOAD_DIR.mkdir(exist_ok=True)

# 2.2. 阿里云百炼 DeepSeek --->50 万 tokens 的 DeepSeek-R1-250120 推理额度（可支撑1000 ~ 5000次简单对话。），有效期2年，用完再按量付费即可
llm = ChatOpenAI(
    openai_api_base="https://ark.cn-beijing.volces.com/api/v3",
    openai_api_key="2f357ee4-84ff-4460-9ba9-ee25d0f9dace",  # 直接填写
    # openai_api_key=os.getenv("ARK_API_KEY"),  # 建议放环境变量
    model="deepseek-r1-250120",
    temperature=0.2,
    max_tokens=2048
)

# -------------------- Prompt 模板 --------------------
prompt = ChatPromptTemplate.from_messages([
    ("system", "你是资深测试架构师，请将下列需求文本转化为符合 Gherkin 语法的测试场景（.feature）。"
               "请以纯文本 Gherkin 格式输出，不要解释。"),
    ("user", "{req_chunk}")
])

chain = prompt | llm


# -------------------- 接口 --------------------
@app.post("/upload", response_class=PlainTextResponse)
async def upload_and_generate(file: UploadFile = File(...)):
    if not file.filename.endswith(".txt"):
        raise HTTPException(400, "只支持 .txt 文件")
    # 保存文件
    file_path = UPLOAD_DIR / file.filename
    with open(file_path, "wb") as f:
        f.write(await file.read())

    # 1. 加载 + 切片
    loader = TextLoader(str(file_path), encoding="utf-8")
    docs = loader.load()
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=2000, chunk_overlap=200)
    chunks = text_splitter.split_documents(docs)

    # 2. 逐段生成 Gherkin 后拼接
    features = []
    for ch in chunks:
        gherkin = chain.invoke({"req_chunk": ch.page_content}).content
        features.append(gherkin.strip())

    return "\n\n".join(features)


# -------------------- 启动 --------------------
if __name__ == "__main__":
    import uvicorn

    uvicorn.run("analysis_app:app", host="172.16.116.189", port=8071, reload=True)
